import bpy
from bpy.props import StringProperty, BoolProperty, EnumProperty, FloatProperty
from bpy_extras.io_utils import ExportHelper
from mathutils import Vector

bl_info = {
    "name": "Export Selected to UE FBX",
    "blender": (3, 0, 0),
    "category": "Import-Export",
}

class EXPORT_OT_selected_ue_fbx_final(bpy.types.Operator, ExportHelper):
    """Export selected objects to FBX with origin at bottom center (UE axes)"""
    bl_idname = "export_scene.selected_ue_fbx_final"
    bl_label = "Export Selected To UE FBX"
    bl_options = {'REGISTER', 'UNDO'}

    filename_ext = ".fbx"
    filter_glob: StringProperty(default="*.fbx", options={'HIDDEN'})

    # UE推荐的轴向设置
    axis_forward: EnumProperty(
        name="Forward",
        items=[('X','X',''),('Y','Y',''),('Z','Z',''),('-X','-X',''),('-Y','-Y',''),('-Z','-Z','')],
        default='-Y',
        description="UE recommended: -Y"
    )
    axis_up: EnumProperty(
        name="Up", 
        items=[('X','X',''),('Y','Y',''),('Z','Z',''),('-X','-X',''),('-Y','-Y',''),('-Z','-Z','')],
        default='Z',
        description="UE recommended: Z"
    )
    
    # 原点位置选项
    origin_position: EnumProperty(
        name="Origin Position",
        items=[
            ('BOTTOM_CENTER', 'Bottom Center', 'Origin at the bottom center of bounding box'),
            ('CENTER', 'Bounding Box Center', 'Origin at the center of bounding box'),
            ('BOTTOM_CORNER', 'Bottom Corner', 'Origin at the bottom corner (min X,Y,Z)')
        ],
        default='BOTTOM_CENTER',
        description="Where to place the origin point"
    )
    
    # 缩放控制选项
    scale_method: EnumProperty(
        name="Scale Method",
        items=[
            ('NONE', 'No Scale (1:1)', 'Export with original scale, no unit conversion'),
            ('UNIT', 'Apply Unit Scale', 'Apply Blender unit scale (may cause size changes)'),
            ('CUSTOM', 'Custom Scale', 'Apply custom scale factor')
        ],
        default='NONE',
        description="How to handle scale during export"
    )
    
    custom_scale: FloatProperty(
        name="Custom Scale Factor",
        default=1.0,
        min=0.001,
        max=1000.0,
        description="Custom scale factor (only used when Scale Method is Custom)"
    )

    def get_origin_point(self, objects):
        """根据选项计算原点位置"""
        if not objects:
            return Vector((0, 0, 0))
        
        all_corners = []
        for obj in objects:
            world_corners = [obj.matrix_world @ Vector(corner) for corner in obj.bound_box]
            all_corners.extend(world_corners)
        
        if not all_corners:
            return Vector((0, 0, 0))
        
        min_v = Vector((min(c.x for c in all_corners),
                        min(c.y for c in all_corners),
                        min(c.z for c in all_corners)))
        max_v = Vector((max(c.x for c in all_corners),
                        max(c.y for c in all_corners),
                        max(c.z for c in all_corners)))
        
        if self.origin_position == 'BOTTOM_CENTER':
            # 水平居中，竖直在底部
            return Vector(((min_v.x + max_v.x) * 0.5,
                          (min_v.y + max_v.y) * 0.5,
                          min_v.z))
        elif self.origin_position == 'CENTER':
            # 完全居中
            return (min_v + max_v) * 0.5
        else:  # BOTTOM_CORNER
            # 底部角落
            return min_v

    def execute(self, context):
        # 获取选中的对象
        sel_objs = [o for o in context.selected_objects if o.type in {'MESH', 'CURVE', 'SURFACE', 'META', 'FONT', 'ARMATURE', 'LATTICE'}]
        if not sel_objs:
            self.report({'ERROR'}, "Please select at least one exportable object.")
            return {'CANCELLED'}

        # 保存场景状态
        original_active = context.view_layer.objects.active
        original_selection = list(context.selected_objects)

        # 计算原点位置
        origin_point = self.get_origin_point(sel_objs)
        self.report({'INFO'}, f"Origin point: ({origin_point.x:.3f}, {origin_point.y:.3f}, {origin_point.z:.3f})")

        # 创建临时Empty作为导出根，放在世界原点
        empty_name = "TEMP_EXPORT_ROOT"
        empty_obj = bpy.data.objects.new(empty_name, None)
        context.collection.objects.link(empty_obj)
        empty_obj.location = (0, 0, 0)
        empty_obj.rotation_euler = (0, 0, 0)
        empty_obj.scale = (1, 1, 1)

        # 如果使用自定义缩放，应用到Empty
        if self.scale_method == 'CUSTOM':
            empty_obj.scale = (self.custom_scale, self.custom_scale, self.custom_scale)

        # 备份原始状态
        backup_data = {}
        for obj in sel_objs:
            backup_data[obj] = {
                'parent': obj.parent,
                'parent_type': obj.parent_type,
                'parent_bone': obj.parent_bone,
                'parent_inverse': obj.matrix_parent_inverse.copy(),
                'matrix_world': obj.matrix_world.copy()
            }

        try:
            # 设置父子关系和重新定位
            for obj in sel_objs:
                # 计算对象当前位置相对于新原点的偏移
                offset = obj.matrix_world.translation - origin_point
                
                # 设置Empty为父对象
                obj.parent = empty_obj
                obj.parent_type = 'OBJECT'
                obj.parent_bone = ""
                
                # 设置对象在Empty坐标系下的位置
                obj.location = offset
                
                # 清空parent_inverse，让Blender自动计算
                obj.matrix_parent_inverse.identity()

            # 选择Empty和所有子对象进行导出
            bpy.ops.object.select_all(action='DESELECT')
            empty_obj.select_set(True)
            for obj in sel_objs:
                obj.select_set(True)
            context.view_layer.objects.active = empty_obj

            # 根据缩放方法设置导出参数
            if self.scale_method == 'NONE':
                apply_unit_scale = False
                scale_options = 'FBX_SCALE_NONE'
            elif self.scale_method == 'UNIT':
                apply_unit_scale = True
                scale_options = 'FBX_SCALE_UNITS'
            else:  # CUSTOM
                apply_unit_scale = False
                scale_options = 'FBX_SCALE_ALL'

            # 执行FBX导出
            bpy.ops.export_scene.fbx(
                filepath=self.filepath,
                use_selection=True,
                apply_unit_scale=apply_unit_scale,
                apply_scale_options=scale_options,
                axis_forward=self.axis_forward,
                axis_up=self.axis_up,
                bake_space_transform=False,
                object_types={'ARMATURE', 'MESH', 'OTHER'},
                use_mesh_modifiers=True,
                add_leaf_bones=False,
                primary_bone_axis='Y',
                secondary_bone_axis='X',
                use_armature_deform_only=True,
                mesh_smooth_type='OFF'
            )

        finally:
            # 恢复对象状态
            for obj in sel_objs:
                if obj.name not in bpy.data.objects:
                    continue
                    
                data = backup_data[obj]
                
                # 先清除父子关系
                obj.parent = None
                obj.matrix_parent_inverse.identity()
                
                # 恢复原始状态
                obj.parent = data['parent']
                obj.parent_type = data['parent_type']
                obj.parent_bone = data['parent_bone']
                obj.matrix_parent_inverse = data['parent_inverse']
                obj.matrix_world = data['matrix_world']

            # 删除临时Empty
            bpy.data.objects.remove(empty_obj, do_unlink=True)

            # 恢复选择状态
            bpy.ops.object.select_all(action='DESELECT')
            for obj in original_selection:
                if obj.name in bpy.data.objects:
                    obj.select_set(True)
            if original_active and original_active.name in bpy.data.objects:
                context.view_layer.objects.active = original_active

        self.report({'INFO'}, f"Successfully exported FBX with {self.origin_position.lower().replace('_', ' ')} origin to: {self.filepath}")
        return {'FINISHED'}

def menu_func(self, context):
    self.layout.operator(EXPORT_OT_selected_ue_fbx_final.bl_idname)

def register():
    bpy.utils.register_class(EXPORT_OT_selected_ue_fbx_final)
    bpy.types.TOPBAR_MT_file_export.append(menu_func)

def unregister():
    bpy.types.TOPBAR_MT_file_export.remove(menu_func)
    bpy.utils.unregister_class(EXPORT_OT_selected_ue_fbx_final)

if __name__ == "__main__":
    register()